from flask import Blueprint, request, jsonify, session, current_app
from logging_config import logger
from flask_jwt_extended import jwt_required, get_jwt_identity, create_access_token
from utils.auth import (
    get_current_user, track_login_attempt, log_auth_event, ROLE_USER,
    revoke_jwt, is_jwt_revoked
)
from utils.encryption import EncryptionUtils
from utils.validation import validate_register_data, validate_login_data
from extensions import db
from models.user_models import User
from flask import current_app

# 延迟导入Redis客户端以避免循环依赖
from flask import current_app

def get_redis_client():
    return current_app.redis_client
import uuid
import datetime
import os

auth_bp = Blueprint('auth', __name__, url_prefix='/api/auth')

@auth_bp.route('/register', methods=['POST'])
def register():
    """用户注册API（增强版）"""
    try:
        logger.info("Register endpoint accessed")
        logger.debug(f"Received registration request: method={request.method}, url={request.url}, content_type={request.content_type}")
        if not request.is_json:
            return jsonify({'error': 'Missing JSON in request'}), 415
        data = request.get_json()
        logger.debug(f"Registration request data: {data}")
        if 'public_key' not in data:
            logger.error("Missing public_key in registration request")
            return jsonify({"status": "error", "message": "public_key is required"}), 400
        # 验证public_key长度
        if len(data['public_key']) > 1000:
            return jsonify({'error': 'public_key exceeds maximum length of 1000 characters'}), 400

        # 验证公钥格式
        try:
            from utils.encryption import EncryptionUtils
            # 尝试加载公钥以验证格式
            public_key_pem = data['public_key'].encode('utf-8')
            EncryptionUtils.load_public_key(public_key_pem)
        except Exception as e:
            return jsonify({'error': f'Invalid public key format: {str(e)}'}), 400

        # 检查公钥唯一性
        from models.user_models import User
        existing_user = User.query.filter_by(public_key=data['public_key']).first()
        if existing_user:
            return jsonify({'error': 'Public key already exists'}), 409
        
        # 输入验证
        validation = validate_register_data(data)
        if validation.get("status") == "error":
            return jsonify(validation), 400
        
        username = data['username']
        password = data['password']
        
        # 检查用户名唯一性（使用参数化查询）
        if User.query.filter_by(username=username).first():
            return jsonify({
                "status": "error",
                "code": 400,
                "errors": {"username": "Username already exists"}
            }), 400
        
        # 生成RSA密钥对

        
        # 生成并安全存储主密钥
        master_key = EncryptionUtils.generate_symmetric_key()  # 256位AES密钥
        kek_salt = os.urandom(16)
        kek = EncryptionUtils.derive_kek(password, kek_salt)
        encrypted_master_key, nonce, tag = EncryptionUtils.aes_encrypt(kek, master_key)
        
        # 保存用户（使用UUID作为主键）
        user_id = str(uuid.uuid4())
        user = User(
            id=user_id,
            username=username,
            public_key=data['public_key'],
            encrypted_master_key=EncryptionUtils.binary_to_hex(encrypted_master_key + nonce + tag),
            kek_salt=EncryptionUtils.binary_to_hex(kek_salt),
            user_role=ROLE_USER,
            is_active=True
        )
        user.set_password(password)  # 假设此方法已处理密码哈希和盐值
        
        db.session.add(user)
        try:
            db.session.commit()
        except Exception as e:
            db.session.rollback()
            logger.error(f"Database error during registration: {str(e)}")
            return jsonify({"status": "error", "message": "Database error occurred"}), 500
        
        # 安全存储私钥（生产环境应使用加密存储服务）

        
        # 生成JWT令牌
        access_token = create_access_token(
            identity=str(user_id),
            additional_claims={"role": ROLE_USER}
        )
        
        log_auth_event(str(user_id), 'REGISTER_SUCCESS', request.remote_addr)
        return jsonify({
            "status": "success",
            "message": "Registration successful",
            "user_id": str(user_id),
            "access_token": access_token
        }), 201
        
    except Exception as e:
        db.session.rollback()
        logger.error(f"Registration error: {str(e)}", exc_info=True)
        return jsonify({
            "status": "error",
            "message": "Registration failed"
        }), 500

@auth_bp.route('/login', methods=['POST'])
def login():
    """用户登录API（增强版）"""
    try:
        data = request.get_json()
        
        # 输入验证
        validation = validate_login_data(data)
        if validation.get("status") == "error":
            return jsonify(validation), 400
        
        username = data['username']
        password = data['password']
        
        logger.info(f"收到登录请求: 用户名={username}")
        
        # 调试信息：打印数据库连接信息
        from flask import current_app
        logger.info(f"当前数据库URI: {current_app.config['SQLALCHEMY_DATABASE_URI']}")
        
        # 尝试直接查询数据库
        from sqlalchemy import text
        try:
            result = db.session.execute(text("SELECT * FROM users WHERE username = :username"), {'username': username})
            raw_user = result.fetchone()
            logger.info(f"原始SQL查询结果: {raw_user}")
        except Exception as e:
            logger.error(f"原始SQL查询出错: {e}")
        
        user = User.query.filter_by(username=username).first()
        logger.info(f"ORM查询结果: {user}")
        
        # 跟踪登录尝试（含频率限制和账户锁定）
        user_id = str(user.id) if user else None
        # 暂时注释掉登录频率限制，便于测试
        # if not track_login_attempt(user_id, request.remote_addr):
        #     return jsonify({
        #         "status": "error",
        #         "code": 429,
        #         "message": "Too many login attempts, please try again later"
        #     }), 429
        
        # 验证用户存在和密码
        logger.info(f"Login attempt for username: {username}")
        logger.info(f"User found: {user is not None}")
        if user:
            logger.info(f"User ID: {user.id}")
            logger.info(f"User active: {user.is_active}")
            logger.info(f"User locked until: {user.login_locked_until}")
            password_match = user.check_password(password)
            logger.info(f"Password match: {password_match}")
            if not password_match:
                logger.warning(f"Password mismatch for user: {username}")
        
        if not user or not user.check_password(password):
            log_auth_event(user_id, 'LOGIN_FAILED', request.remote_addr)
            return jsonify({
                "status": "error",
                "code": 401,
                "message": "Invalid credentials"
            }), 401
        
        # 检查账号是否活跃或锁定
        if not user.is_active:
            return jsonify({
                "status": "error",
                "code": 403,
                "message": "Account disabled"
            }), 403
        
        # 暂时注释掉账户锁定检查，便于测试
        # 修复时间戳比较问题
        # from datetime import timezone
        # now = datetime.datetime.now(timezone.utc)
        # if user.login_locked_until:
        #     # 确保login_locked_until也是时区感知的
        #     locked_until = user.login_locked_until
        #     if locked_until.tzinfo is None:
        #         # 如果login_locked_until是naive datetime，转换为UTC
        #         locked_until = locked_until.replace(tzinfo=timezone.utc)
        #     
        #     if locked_until > now:
        #         lock_seconds = (locked_until - now).total_seconds()
        #         return jsonify({
        #             "status": "error",
        #             "code": 423,
        #             "message": "Account locked",
        #             "lock_seconds": int(lock_seconds)
        #         }), 423
        
        # 解密主密钥
        try:
            kek_salt = EncryptionUtils.hex_to_binary(user.kek_salt)
            logger.debug(f"kek_salt length: {len(kek_salt)} bytes")
            
            kek = EncryptionUtils.derive_kek(password, kek_salt)
            logger.debug(f"kek length: {len(kek)} bytes")
            
            encrypted_master_key = EncryptionUtils.hex_to_binary(user.encrypted_master_key)
            logger.debug(f"encrypted_master_key length: {len(encrypted_master_key)} bytes")
            
            # 验证总长度是否足够
            if len(encrypted_master_key) < 28:
                raise ValueError(f"encrypted_master_key too short: {len(encrypted_master_key)} bytes")
            
            ciphertext = encrypted_master_key[:-28]
            nonce = encrypted_master_key[-28:-16]
            tag = encrypted_master_key[-16:]
            
            logger.debug(f"ciphertext length: {len(ciphertext)} bytes")
            logger.debug(f"nonce length: {len(nonce)} bytes")
            logger.debug(f"tag length: {len(tag)} bytes")
            
            # 验证nonce和tag长度
            if len(nonce) != 12:
                raise ValueError(f"Invalid nonce length: {len(nonce)} bytes (expected 12)")
            if len(tag) != 16:
                raise ValueError(f"Invalid tag length: {len(tag)} bytes (expected 16)")
            
            master_key = EncryptionUtils.aes_decrypt(kek, ciphertext, nonce, tag)
            logger.debug(f"master_key length: {len(master_key)} bytes")
        except Exception as e:
            logger.error(f"Decryption error: {str(e)}")
            raise
        
        # 安全存储会话数据（使用加密会话或短生命周期令牌）
        session['user_id'] = str(user.id)
        session['master_key_hex'] = EncryptionUtils.binary_to_hex(master_key)
        session.permanent = True  # 设置会话过期时间（通过配置）
        
        # 更新在线状态
        user.online = True
        user.last_login = datetime.datetime.now(datetime.timezone.utc)
        db.session.commit()
        
        # 生成JWT令牌（带短有效期）
        access_token = create_access_token(
            identity=str(user.id),
            expires_delta=datetime.timedelta(hours=1),
            additional_claims={"role": user.user_role}  # 使用正确的属性名
        )
        
        log_auth_event(str(user.id), 'LOGIN_SUCCESS', request.remote_addr)
        return jsonify({
            "status": "success",
            "message": "Login successful",
            "user_id": str(user.id),
            "access_token": access_token
        }), 200
        
    except Exception as e:
        logger.error(f"Login error: {str(e)}", exc_info=True)
        return jsonify({
            "status": "error",
            "message": "Login failed"
        }), 500

@auth_bp.route('/logout', methods=['POST'])
@jwt_required()
def logout():
    """用户登出API（增强版）"""
    try:
        user_id = get_jwt_identity()
        user = User.query.get(user_id)
        
        # 使当前JWT令牌失效
        token = request.headers.get('Authorization', '').split(' ')[1]
        if token and revoke_jwt(token):
            logger.info(f"JWT revoked: user={user_id}")
        
        if user:
            user.online = False
            db.session.commit()
        
        # 清除会话数据
        session.clear()
        log_auth_event(user_id, 'LOGOUT', request.remote_addr)
        
        return jsonify({
            "status": "success",
            "message": "Logged out successfully"
        }), 200
    except Exception as e:
        logger.error(f"Logout error: {str(e)}", exc_info=True)
        return jsonify({
            "status": "error",
            "message": "Logout failed"
        }), 500

@auth_bp.route('/me', methods=['GET'])
@jwt_required()
def get_current_user_info():
    """获取当前用户信息（增强版）"""
    try:
        user_id = get_jwt_identity()
        user = User.query.get(user_id)
        
        if not user:
            return jsonify({
                "status": "error",
                "code": 404,
                "message": "User not found"
            }), 404
        
        # 检查JWT是否已被吊销
        auth_header = request.headers.get('Authorization', '')
        if auth_header and ' ' in auth_header:
            token = auth_header.split(' ')[1]
            if is_jwt_revoked(token):
                return jsonify({
                    "status": "error",
                    "code": 401,
                    "message": "Token revoked, please login again"
                }), 401
        
        return jsonify({
            "status": "success",
            "data": {
                "id": str(user.id),
                "username": user.username,
                "role": user.user_role,  # 使用正确的属性名
                "online": user.online,
                "created_at": user.created_at.isoformat(),
                "last_login": user.last_login.isoformat() if user.last_login else None
            }
        }), 200
        
    except Exception as e:
        logger.error(f"Get user info error: {str(e)}", exc_info=True)
        return jsonify({
            "status": "error",
            "message": "Failed to get user info"
        }), 500